package com.company.socks5;

import io.netty.bootstrap.Bootstrap;
import io.netty.bootstrap.ServerBootstrap;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufUtil;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.ByteToMessageDecoder;
import io.netty.handler.codec.socksx.v5.*;
import io.netty.handler.logging.LogLevel;
import io.netty.handler.logging.LoggingHandler;
import io.netty.util.Attribute;
import io.netty.util.AttributeKey;

import java.util.List;

import static xyz.noark.log.LogHelper.logger;

public class SocksProxyServer {
    private static final int port = 1080;

    public static void main(String[] args) throws InterruptedException {
        new SocksProxyServer().start();
    }

    public void start() throws InterruptedException {
        logger.info("Starting SOCKS5 proxy on port {}", port);

        EventLoopGroup bossGroup = new NioEventLoopGroup(1);
        EventLoopGroup workerGroup = new NioEventLoopGroup();
        try {
            new ServerBootstrap().group(bossGroup, workerGroup).channel(NioServerSocketChannel.class).handler(new LoggingHandler(LogLevel.INFO)) // 添加服务端启动日志
                    .childHandler(new ChannelInitializer<Channel>() {
                        @Override
                        protected void initChannel(Channel ch) {
                            ch.pipeline().addLast(
                                    //new PacketLoggerHandler(),  // 添加封包日志
                                    //new IdRewriterHandler(), // 添加ID重写处理器
                                    new Socks5InitialRequestDecoder(),  // 添加初始请求处理
                                    Socks5ServerEncoder.DEFAULT,
                                    //new Socks5AuthMethodRequestDecoder(), // 认证阶段
                                    new SimpleChannelInboundHandler<Socks5InitialRequest>() {
                                        @Override
                                        protected void channelRead0(ChannelHandlerContext ctx, Socks5InitialRequest request) {
                                            ctx.pipeline().addLast(new Socks5CommandRequestDecoder(), new Socks5CommandRequestHandler() // 自定义命令处理器
                                            );
                                            ctx.writeAndFlush(new DefaultSocks5InitialResponse(Socks5AuthMethod.NO_AUTH));
                                        }
                                    });
                        }
                    }).bind(port).addListener((ChannelFutureListener) future -> {
                        if (future.isSuccess()) {
                            logger.info("Server started on port {}", port);
                        } else {
                            logger.error("Failed to bind port {}", port, future.cause());
                        }
                    }).sync().channel().closeFuture().sync();
        } finally {
            bossGroup.shutdownGracefully();
            workerGroup.shutdownGracefully();
        }
    }

    private static class Socks5CommandRequestHandler extends SimpleChannelInboundHandler<Socks5CommandRequest> {

        @Override
        protected void channelRead0(ChannelHandlerContext ctx, Socks5CommandRequest request) {
            if (request.type() == Socks5CommandType.CONNECT) {
                Bootstrap b = new Bootstrap();
                b.group(ctx.channel().eventLoop()).channel(NioSocketChannel.class).handler(new ChannelInitializer<Channel>() {
                    @Override
                    protected void initChannel(Channel ch) {
                        ch.pipeline().addLast(new RelayHandler(ctx.channel(), false) // 双向转发
                        );
                    }
                });

                b.connect(request.dstAddr(), request.dstPort()).addListener((ChannelFuture future) -> {
                    if (future.isSuccess()) {
                        ctx.pipeline().addLast(new SocksProxyFrameDecoder(), new RelayHandler(future.channel(), true));
                        ctx.writeAndFlush(new DefaultSocks5CommandResponse(Socks5CommandStatus.SUCCESS, request.dstAddrType(), request.dstAddr(), request.dstPort()));
                    } else {
                        ctx.writeAndFlush(new DefaultSocks5CommandResponse(Socks5CommandStatus.FAILURE, request.dstAddrType()));
                        ctx.close();
                    }
                });
            }
        }
    }

    public static class SocksProxyFrameDecoder extends ByteToMessageDecoder {
        @Override
        protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) {
            while (in.readableBytes() >= 10) {
                int packetLength = in.getIntLE(in.readerIndex());
                if (in.readableBytes() < packetLength) {
                    return; // 等待更多数据
                }
                out.add(in.readRetainedSlice(packetLength));
            }
        }
    }

    private static class RelayHandler extends ChannelInboundHandlerAdapter {
        private final Channel relayChannel;
        private final boolean isClientSend;
        private static final AttributeKey<Integer> ID_COUNTER = AttributeKey.newInstance("IdCounter");

        public RelayHandler(Channel relayChannel, boolean isClientSend) {
            this.relayChannel = relayChannel;
            this.isClientSend = isClientSend;
        }

        @Override
        public void channelRead(ChannelHandlerContext ctx, Object msg) {
            if (isClientSend && msg instanceof ByteBuf buf) {
                Attribute<Integer> attr = ctx.channel().attr(ID_COUNTER);
                Integer id = attr.get();
                if (id == null) {
                    id = 100;
                }
                attr.set(id + 1);


                if (buf.readableBytes() >= 10) { // 至少包含长度+ID
                    int packetLength = buf.getIntLE(0);
                    if (buf.readableBytes() == packetLength) {
                        buf.setIntLE(6, id); // 修改ID字段
                        logger.debug("Rewrote ID to: {}", id);
                    }
                }

                logger.info("buf2={}", ByteBufUtil.prettyHexDump(buf));
            }

            relayChannel.writeAndFlush(msg);
        }

        @Override
        public void channelInactive(ChannelHandlerContext ctx) {
            relayChannel.close();
        }
    }
}
